#!/usr/bin/python3

import argparse
import copy
import functools
import logging
import subprocess
import os
import sys
import stat

import fedora_messaging.api
import fedora_messaging.config
import fedora_messaging.exceptions

from pathlib import Path

logging.basicConfig(level=logging.ERROR)
logger = logging.getLogger('updates-sync')


SOURCE = '/mnt/koji/compose/updates/'
FEDORADEST = '/pub/fedora/linux/updates/'
FEDORAMODDEST = '/pub/fedora/linux/modular/updates/'
FEDORAALTDEST = '/pub/fedora-secondary/updates/'
EPELDEST = '/pub/epel/'
OSTREESOURCE = '/mnt/koji/compose/ostree/repo/'
OSTREEDEST = '/mnt/koji/ostree/repo/'
RELEASES = {'f37': {'topic': 'fedora',
                    'version': '37',
                    'modules': ['fedora', 'fedora-secondary'],
                    'repos': {'updates': {
                        'from': 'f37-updates',
                        'ostrees': [{'ref': 'fedora/37/%(arch)s/updates/silverblue',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']},
                                    {'ref': 'fedora/37/%(arch)s/updates/kinoite',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']}],
                        'to': [{'arches': ['x86_64', 'armhfp', 'aarch64', 'source'],
                                'dest': os.path.join(FEDORADEST, '37', 'Everything')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, '37', 'Everything')}
                              ]},
                              'updates-testing': {
                        'from': 'f37-updates-testing',
                        'ostrees': [{'ref': 'fedora/37/%(arch)s/testing/silverblue',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']},
                                    {'ref': 'fedora/37/%(arch)s/testing/kinoite',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']}],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, 'testing', '37', 'Everything')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, 'testing', '37', 'Everything')}
                              ]}}
                   },
            'f37m': {'topic': 'fedora',
                    'version': '37m',
                    'modules': ['fedora', 'fedora-secondary'],
                    'repos': {'updates': {
                        'from': 'f37-modular-updates',
                        'ostrees': [],
                        'to': [{'arches': ['x86_64', 'aarch64', 'source'],
                                'dest': os.path.join(FEDORADEST, '37', 'Modular')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, '37', 'Modular')}
                              ]},
                              'updates-testing': {
                        'from': 'f37-modular-updates-testing',
                        'ostrees': [],
                        'to': [{'arches': ['x86_64', 'aarch64', 'source'],
                                'dest': os.path.join(FEDORADEST, 'testing', '37', 'Modular')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, 'testing', '37', 'Modular')}
                              ]}}
                   },
            'f36': {'topic': 'fedora',
                    'version': '36',
                    'modules': ['fedora', 'fedora-secondary'],
                    'repos': {'updates': {
                        'from': 'f36-updates',
                        'ostrees': [{'ref': 'fedora/36/%(arch)s/updates/silverblue',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']},
                                    {'ref': 'fedora/36/%(arch)s/updates/kinoite',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']}],
                        'to': [{'arches': ['x86_64', 'armhfp', 'aarch64', 'source'],
                                'dest': os.path.join(FEDORADEST, '36', 'Everything')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, '36', 'Everything')}
                              ]},
                              'updates-testing': {
                        'from': 'f36-updates-testing',
                        'ostrees': [{'ref': 'fedora/36/%(arch)s/testing/silverblue',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']},
                                    {'ref': 'fedora/36/%(arch)s/testing/kinoite',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']}],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, 'testing', '36', 'Everything')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, 'testing', '36', 'Everything')}
                              ]}}
                   },
            'f36m': {'topic': 'fedora',
                    'version': '36m',
                    'modules': ['fedora', 'fedora-secondary'],
                    'repos': {'updates': {
                        'from': 'f36-modular-updates',
                        'ostrees': [],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, '36', 'Modular')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, '36', 'Modular')}
                              ]},
                              'updates-testing': {
                        'from': 'f36-modular-updates-testing',
                        'ostrees': [],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, 'testing', '36', 'Modular')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, 'testing', '36', 'Modular')}
                              ]}}
                   },
            'f35': {'topic': 'fedora',
                    'version': '35',
                    'modules': ['fedora', 'fedora-secondary'],
                    'repos': {'updates': {
                        'from': 'f35-updates',
                        'ostrees': [{'ref': 'fedora/35/%(arch)s/updates/silverblue',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']},
                                    {'ref': 'fedora/35/%(arch)s/updates/kinoite',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']}],
                        'to': [{'arches': ['x86_64', 'armhfp', 'aarch64', 'source'],
                                'dest': os.path.join(FEDORADEST, '35', 'Everything')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, '35', 'Everything')}
                              ]},
                              'updates-testing': {
                        'from': 'f35-updates-testing',
                        'ostrees': [{'ref': 'fedora/35/%(arch)s/testing/silverblue',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']},
                                    {'ref': 'fedora/35/%(arch)s/testing/kinoite',
                                     'dest': OSTREEDEST,
                                     'arches': ['x86_64', 'ppc64le', 'aarch64']}],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, 'testing', '35', 'Everything')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, 'testing', '35', 'Everything')}
                              ]}}
                   },
            'f35m': {'topic': 'fedora',
                    'version': '35m',
                    'modules': ['fedora', 'fedora-secondary'],
                    'repos': {'updates': {
                        'from': 'f35-modular-updates',
                        'ostrees': [],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, '35', 'Modular')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, '35', 'Modular')}
                              ]},
                              'updates-testing': {
                        'from': 'f35-modular-updates-testing',
                        'ostrees': [],
                        'to': [{'arches': ['x86_64', 'aarch64', 'armhfp', 'source'],
                                'dest': os.path.join(FEDORADEST, 'testing', '35', 'Modular')},
                               {'arches': ['ppc64le', 's390x'],
                                'dest': os.path.join(FEDORAALTDEST, 'testing', '35', 'Modular')}
                              ]}}
                   },
            'epel9': {'topic': 'epel',
                      'version': '9',
                      'modules': ['epel'],
                      'repos': {'epel-testing': {
                          'from': 'epel9-testing',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                  'dest': os.path.join(EPELDEST, 'testing', '9', 'Everything')}
                                ]},
                                'epel': {
                          'from': 'epel9',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                  'dest': os.path.join(EPELDEST, '9', 'Everything')}
                                ]}}
                     },
            'epel9n': {'topic': 'epel',
                      'version': '9',
                      'modules': ['epel'],
                      'repos': {'epel-testing': {
                          'from': 'epel9-next-testing',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                  'dest': os.path.join(EPELDEST, 'testing', 'next', '9', 'Everything')}
                                ]},
                                'epel': {
                          'from': 'epel9-next',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                  'dest': os.path.join(EPELDEST, 'next', '9', 'Everything')}
                                ]}}
                     },
            'epel8': {'topic': 'epel',
                      'version': '8',
                      'modules': ['epel'],
                      'repos': {'epel-testing': {
                          'from': 'epel8-testing',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                  'dest': os.path.join(EPELDEST, 'testing', '8', 'Everything')}
                                ]},
                                'epel': {
                          'from': 'epel8',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                  'dest': os.path.join(EPELDEST, '8', 'Everything')}
                                ]}}
                     },
            'epel8n': {'topic': 'epel',
                      'version': '8',
                      'modules': ['epel'],
                      'repos': {'epel-testing': {
                          'from': 'epel8-next-testing',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 'source'],
                                  'dest': os.path.join(EPELDEST, 'testing', 'next', '8', 'Everything')}
                                ]},
                                'epel': {
                          'from': 'epel8-next',
                          'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 'source'],
                                  'dest': os.path.join(EPELDEST, 'next', '8', 'Everything')}
                                ]}}
                     },
            'epel8m': {'topic': 'epel',
                    'version': '8m',
                    'modules': ['epel'],
                    'repos': {'epel': {
                        'from': 'epel8-modular-updates',
                        'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                'dest': os.path.join(EPELDEST, '8', 'Modular')},
                              ]},
                              'epel-testing': {
                        'from': 'epel8-modular-updates-testing',
                        'to': [{'arches': ['x86_64', 'aarch64', 'ppc64le', 's390x', 'source'],
                                'dest': os.path.join(EPELDEST, 'testing', '8', 'Modular')},
                              ]}}
                   },
            'epel7': {'topic': 'epel',
                      'version': '7',
                      'modules': ['epel'],
                      'repos': {'epel-testing': {
                          'from': 'epel7-testing',
                          'to': [{'arches': ['x86_64', 'ppc64le', 'source'],
                                  'dest': os.path.join(EPELDEST, 'testing', '7')}
                                ]},
                                'epel': {
                          'from': 'epel7',
                          'to': [{'arches': ['x86_64', 'ppc64le', 'source'],
                                  'dest': os.path.join(EPELDEST, '7')}
                                ]}}
                     },
           }


# Beneath this is code, no config needed here

def run_command(cmd):
    logger.info('Running %s', cmd)
    return subprocess.check_output(cmd,
                                   stderr=subprocess.STDOUT,
                                   shell=False)


def get_ostree_ref(repo, ref):
    reffile = os.path.join(repo, 'refs', 'heads', ref)
    if not os.path.exists(reffile):
        return '----'
    with open(reffile, 'r') as f:
        return f.read().split()[0]


def sync_ostree(dst, ref):
    src_commit = get_ostree_ref(OSTREESOURCE, ref)
    dst_commit = get_ostree_ref(dst, ref)
    if src_commit == dst_commit:
        logger.info('OSTree at %s, ref %s in sync', dst, ref)
    else:
        # Set the umask to be more permissive so directories get group write
        # https://github.com/ostreedev/ostree/pull/1984
        oldumask = os.umask(0o0002)
        try:
            print('Syncing ostree ref %s: %s -> %s'
                  % (ref, src_commit, dst_commit))
            logger.info('Syncing OSTree to %s, ref %s: %s -> %s',
                        dst, ref, src_commit, dst_commit)
            cmd = ['ostree', 'pull-local', '--verbose', '--repo',
                   dst, OSTREESOURCE, ref]
            out = run_command(cmd)
            cmd = ['ostree', 'summary', '--verbose', '--repo', dst, '--update']
            run_command(cmd)
            print('Ostree ref %s now at %s' % (ref, src_commit))
        finally:
            os.umask(oldumask)


def update_fullfilelist(modules):
    if not modules:
        logger.info('No filelists to update')
        return
    cmd = ['/usr/local/bin/update-fullfiletimelist', '-l',
           '/pub/fedora-secondary/update-fullfiletimelist.lock', '-t', '/pub']
    cmd.extend(modules)
    run_command(cmd)


def rsync(from_path, to_path, excludes=[], link_dest=None, delete=False):
    cmd = ['rsync', '-rlptDvHh', '--stats', '--no-human-readable']
    for excl in excludes:
        cmd += ['--exclude', excl]
    if link_dest:
        cmd += ['--link-dest', link_dest]
    if delete:
        cmd += ['--delete', '--delete-delay']

    cmd += [from_path, to_path]

    stdout = run_command(cmd)

    results = {'num_bytes': 0,
               'num_deleted': 0}
    for line in stdout.decode('utf-8').split('\n'):
        if 'Literal data' in line:
            results['num_bytes'] = int(line.split()[2])
        elif 'deleting ' in line:
            results['num_deleted'] += 1

    return results


def collect_stats(stats):
    to_collect = ['num_bytes', 'num_deleted']
    totals = {}
    for stat in to_collect:
        totals[stat] = functools.reduce(lambda x, y: x + y,
                                        [r[stat] for r in stats])
    return totals


def to_human(num_bytes):
    ranges = ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB', 'ZiB', 'YiB']
    cur_range = 0
    while num_bytes >= 1024 and cur_range+1 <= len(ranges):
        num_bytes = num_bytes / 1024
        cur_range += 1
    return '%s %s' % (num_bytes, ranges[cur_range])



def sync_single_repo_arch(release, repo, arch, dest_path):
    source_path = os.path.join(SOURCE,
                               RELEASES[release]['repos'][repo]['from'],
                               'compose', 'Everything', arch)

    maindir = 'tree' if arch == 'source' else 'os'

    results = []
    do_drpms = os.path.exists(os.path.join(source_path, maindir, 'drpms'))

    results.append(rsync(os.path.join(source_path, maindir, 'Packages'),
                         os.path.join(dest_path)))
    if do_drpms:
        results.append(rsync(os.path.join(source_path, maindir, 'drpms'),
                             os.path.join(dest_path)))
    if arch != 'source':
        results.append(rsync(os.path.join(source_path, 'debug', 'tree', 'Packages'),
                             os.path.join(dest_path, 'debug')))
        results.append(rsync(os.path.join(source_path, 'debug', 'tree', 'repodata'),
                             os.path.join(dest_path, 'debug'),
                             delete=True))
        results.append(rsync(os.path.join(source_path, 'debug', 'tree', 'Packages'),
                             os.path.join(dest_path, 'debug'),
                             delete=True))
    results.append(rsync(os.path.join(source_path, maindir, 'repodata'),
                         os.path.join(dest_path),
                         delete=True))
    results.append(rsync(os.path.join(source_path, maindir, 'Packages'),
                         os.path.join(dest_path),
                         delete=True))
    if do_drpms:
        results.append(rsync(os.path.join(source_path, maindir, 'drpms'),
                             os.path.join(dest_path),
                             delete=True))

    return collect_stats(results)


def sync_single_repo(release, repo):
    global FEDMSG_INITED

    results = []

    for archdef in RELEASES[release]['repos'][repo]['to']:
        for arch in archdef['arches']:
            if 'source' in arch:
                dest_path = os.path.join(archdef['dest'], arch, 'tree')
            else:
                dest_path = os.path.join(archdef['dest'], arch)
            results.append(sync_single_repo_arch(release, repo, arch, dest_path))

    stats = collect_stats(results)
    fedmsg_msg = {'repo': repo,
                  'release': RELEASES[release]['version'],
                  'bytes': to_human(stats['num_bytes']),
                  'raw_bytes': str(stats['num_bytes']),
                  'deleted': str(stats['num_deleted'])}

    try:
        msg = fedora_messaging.api.Message(
            topic="bodhi.updates.{}.sync".format(RELEASES[release]['topic']),
            body=fedmsg_msg
        )
        fedora_messaging.api.publish(msg)
    except fedora_messaging.exceptions.PublishReturned as e:
        print("Fedora Messaging broker rejected message %s: %s" % (msg.id, e))
    except fedora_messaging.exceptions.ConnectionException as e:
        print("Error sending message %s: %s" % (msg.id, e))
    except Exception as e:
        print("Error sending fedora-messaging message: %s" % (e))


def determine_last_link(release, repo):
    source_path = os.path.join(SOURCE,
                               RELEASES[release]['repos'][repo]['from'])
    target = os.readlink(source_path)
    logger.info('Release %s, repo %s, target %s', release, repo, target)
    RELEASES[release]['repos'][repo]['from'] = target
    return target


def sync_single_release(release):
    needssync = False

    for repo in RELEASES[release]['repos']:
        target = determine_last_link(release, repo)

        # if "to" is not empty then sync repo
        if RELEASES[release]['repos'][repo]['to']:
            curstatefile = os.path.join(
                RELEASES[release]['repos'][repo]['to'][0]['dest'], 'state')
            curstate = None
            if os.path.exists(curstatefile):
                with open(curstatefile, 'r') as f:
                    curstate = f.read().split()[0]

            # Resync if Bodhi failed out during the sync waiting, which leads
            # to changed repomd.xml without an updated repo.
            # (updateinfo is inserted again)
            # Fix: https://github.com/fedora-infra/bodhi/pull/1986
            if curstate and curstate == target:
                curstatestat = os.stat(curstatefile)
                repostat = os.stat(os.path.join(
                    target, 'compose', 'Everything',
                    RELEASES[release]['repos'][repo]['to'][0]['arches'][0],
                    'os', 'repodata', 'repomd.xml'))
                if curstatestat[stat.ST_MTIME] < repostat[stat.ST_MTIME]:
                    # If the curstate file has an earlier mtime than the repomd
                    # of the first architecture, this repo was re-generated
                    # after the first time it got staged. Resync.
                    logger.error(
                        'Re-stage detected of %s %s. '
                        'State mtime: %s, repo mtime: %s',
                        release, repo,
                        curstatestat[stat.ST_MTIME],
                        repostat[stat.ST_MTIME])
                    curstate = None

            if curstate and curstate == target:
                logger.info('This repo has already been synced')
            else:
                print('Syncing %s %s from %s -> %s' % (release,
                                                       repo,
                                                       curstate,
                                                       target))
                sync_single_repo(release, repo)
                with open(curstatefile, 'w') as f:
                    f.write(target)
                needssync = True
                print('Synced %s %s to %s' % (release, repo, target))

        for ostree in RELEASES[release]['repos'][repo].get('ostrees', []):
            pairs = []
            for arch in ostree.get('arches', ['']):
                dest = ostree['dest'] % {'arch': arch}
                ref = ostree['ref'] % {'arch': arch}
                pairs.append((dest, ref))

            for pair in pairs:
                sync_ostree(*pair)

    return needssync

def get_epel_release_rel_path(release, epel_next=False):
    """
    Get the relative path of the epel-{next-}release build
    """
    if epel_next:
        for path in Path(RELEASES[release]['repos']['epel']['to'][0]['dest']).rglob('epel-next-release*noarch*'):
            if 'Packages' in str(path) and 'x86_64' in str(path):
                pkg_relpath = os.path.relpath(path,EPELDEST)
    else:
        for path in Path(RELEASES[release]['repos']['epel']['to'][0]['dest']).rglob('epel-release*noarch*'):
            if 'Packages' in str(path) and 'x86_64' in str(path):
                pkg_relpath = os.path.relpath(path,EPELDEST)
    return path, pkg_relpath

def update_epel_release_latest(releases):
    """
    This function, creates or updates a symbolic links for epel-release, latest and next, packages.

    Creates or updates a symbolic link pointing to the latest release of the epel-release package and 
    another pointing to the next release of the epel-release package.
    
    The symbolic link will be created or updated if:
    - There isn't a symbolic link for the latest package;
    - Current symbolic link is pointing to an outdated package;
    - Current symbolic link is broken;
    - There is a file that isn't a link with the same name of the symbolic link.

    If the symbolic link is pointing to the latest release already, this function will do nothing.

    Parameters:
    releases (dict): contains similar information of global variable RELEASES
    """
    for release in releases:
        if 'epel' in release and 'Modular' not in RELEASES[release]['repos']['epel']['to'][0]['dest']:
            if 'next' in RELEASES[release]['repos']['epel']['to'][0]['dest']:
                dest = '/pub/epel/epel-next-release-latest-' + release[4] + '.noarch.rpm'
                # For next's epel release, use the subpackage rpm from epel repo instead of
                # epel next repo
                release = release[:-1]
                path, pkg_relpath = get_epel_release_rel_path(release, True)
            else:
                dest = '/pub/epel/epel-release-latest-' + release[4] + '.noarch.rpm'
                path, pkg_relpath = get_epel_release_rel_path(release)

            if os.path.lexists(dest) and os.path.islink(dest):
                origin_dest = os.path.join(EPELDEST,os.readlink(dest))
                if origin_dest != str(path):
                    os.remove(dest)
                    os.symlink(pkg_relpath, dest)
            elif os.path.lexists(dest) and not os.path.islink(dest):
                os.remove(dest)
                os.symlink(pkg_relpath, dest)
            else:
                os.symlink(pkg_relpath, dest)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("releases", nargs="*", default=RELEASES)
    parser.add_argument(
        "--config",
        dest="config",
        help="fedora-messaging configuration file to use. "
        "This allows overriding the default "
        "/etc/fedora-messaging/config.toml.",
    )
    args = parser.parse_args()

    if args.config:
        fedora_messaging.config.conf.load_config(args.config)

    to_update = []

    for release in args.releases:
        if sync_single_release(release):
            to_update.extend(RELEASES[release]['modules'])

    to_update = list(set(to_update))
    logger.info('Filelists to update: %s', to_update)
    update_fullfilelist(to_update)
    update_epel_release_latest(args.releases)

if __name__ == '__main__':
    main()
